#ifndef _ZQ_CPP_LIB_GRADIENT_COMPRESSION_BODY_ONEBIT_BODY_H
#define _ZQ_CPP_LIB_GRADIENT_COMPRESSION_BODY_ONEBIT_BODY_H
#include <stdint.h>
#include <thrust/random.h>
#include <thrust/sort.h>                       //sort()
#include <thrust/execution_policy.h>           //thrust::device
#include <thrust/functional.h>                 //greater<float>
#include <thrust/copy.h>                       //copy_if
#include <thrust/iterator/counting_iterator.h> // counting_iterator
#include <thrust/transform.h>                  //trnasform
#include "../naive_random.hpp"
#include "../operate_memory/get_policy_general.h"
#include "../likely.h"
#include <chrono>

namespace zq_cpp_lib
{
namespace gradient_compression_body
{
using namespace zq_cpp_lib::operate_memory;

struct abs_add{
    __host__ __device__
    float operator()(const float& x, const float& y){
        return abs(x)+abs(y);
    }
};

struct compress{
    float* in_float;
    int in_float_size;
    compress(
        float* in_float_,
        int in_float_size_
    ){
        in_float = in_float_;
        in_float_size = in_float_size_;
    };
    __host__ __device__
    uint8_t operator()(const int32_t& i){
        int j;
        uint8_t q= in_float[i*8] < 0; // first float should always valid.
        // uint8_t q = 0;
        int t;
        // q:       _|_|_|_|_|_|_|_ 
        // j=0      0
        // j=1        1
        // j=7                    7
        // (q>>(7-j)) & 1
        for (j = 1; j < 8; j++){
            t = i*8+j;
            q<<=1;
            // if (t>1000) printf("t=%d\n", t);
            if (t < in_float_size){
                // if (t>1000) printf("t=%d\n", t);
                q = q | (in_float[t]<0);
            }
        }
        return q;
    }
};

struct decompress_add_to{
    uint8_t* in_uint8_t;
    float scale;
    float* out_float;
    decompress_add_to(
        uint8_t* in_uint8_t_,
        float scale_,
        float* out_float_
    ){
        in_uint8_t = in_uint8_t_;
        scale = scale_;
        out_float = out_float_;
    }
    __host__ __device__
    float operator()(const int32_t& i){
        uint8_t j = i&7;
        uint8_t q = (in_uint8_t[i/8]>>(7-j))&1;
        return out_float[i] + scale*(1-2*q);
    }
};

struct decompress_write_to{
    uint8_t* in_uint8_t;
    float scale;
    float* out_float;
    decompress_write_to(
        uint8_t* in_uint8_t_,
        float scale_,
        float* out_float_
    ){
        in_uint8_t = in_uint8_t_;
        scale = scale_;
        out_float = out_float_;
    }
    __host__ __device__
    float operator()(const int32_t& i){
        uint8_t j = i&7;
        uint8_t q = (in_uint8_t[i/8]>>(7-j))&1;
        return scale*(1-2*q);
    }
};



template<typename policy_t>
int onebit_body(
    float* in_float,
    int32_t in_float_size,
    uint8_t* out_uint8_t,
    int32_t out_uint8_t_size, // only for checking enough space
    uint8_t use_scalar,
    policy_t policy,
    void* stream
){
    // 0 --> 0
    // 1~8 --> 1
    // 9~16 --> 2
    int32_t packed_size = (in_float_size-1)/8+1;
    if (packed_size+4>out_uint8_t_size){
        return -1;
    }

    float scale = 1.0f;
    if (use_scalar && in_float_size){
        float sum = thrust::reduce(
            policy,
            in_float,
            in_float+in_float_size,
            0.0f,
            abs_add()
        );
        scale = sum / in_float_size;
    };
    
    thrust::counting_iterator<int32_t> index_sequence_begin(0);
    thrust::transform(
        policy,
        index_sequence_begin,
        index_sequence_begin + packed_size,
        out_uint8_t,
        compress(
            in_float,
            in_float_size
        )
    );
    get_policy<policy_t>::memcpyIn(out_uint8_t+packed_size, &scale, sizeof(float), stream);
    return 0;
}

template<typename policy_t> 
int onebit_r_body(
    uint8_t* in_uint8_t,
    int32_t in_uint8_t_size,
    float* out_float,
    int32_t out_float_size,
    uint8_t use_scalar,
    int32_t original_len, // # elements in original data
    uint8_t is_add_to,
    policy_t policy,
    void* stream
){
    int32_t packed_size = (original_len-1)/8+1;
    float scale = 1.0f;
    if (use_scalar){
        get_policy<policy_t>::memcpyOut(&scale, in_uint8_t+packed_size, sizeof(float), stream);
    }
    thrust::counting_iterator<int32_t> index_sequence_begin(0);
    if (is_add_to){
        thrust::transform(
            policy,
            index_sequence_begin,
            index_sequence_begin+original_len,
            out_float,
            decompress_add_to(
                in_uint8_t,
                scale,
                out_float
            )
        );
    }
    else{
        thrust::transform(
            policy,
            index_sequence_begin,
            index_sequence_begin+original_len,
            out_float,
            decompress_write_to(
                in_uint8_t,
                scale,
                out_float
            )
        );
    }
    return 0;
}

}
}

#endif
